use std::fmt;
use std::str::Chars;

// 定义令牌类型
#[derive(Debug, Clone, PartialEq)]
enum Token {
  Number(f64),
  Plus,     // +
  Minus,    // -
  Star,     // *
  Slash,    // /
  LParen,   // (
  RParen,   // )
  LBracket, // 数组左括号 [
  RBracket, // 数组右括号 ]
  Comma,    // 逗号 ,
  Eof,
}

impl fmt::Display for Token {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    match self {
      Token::Number(n) => write!(f, "Number({})", n),
      Token::Plus => write!(f, "Plus"),
      Token::Minus => write!(f, "Minus"),
      Token::Star => write!(f, "Star"),
      Token::Slash => write!(f, "Slash"),
      Token::LParen => write!(f, "LParen"),
      Token::RParen => write!(f, "RParen"),
      Token::LBracket => write!(f, "LBracket"),
      Token::RBracket => write!(f, "RBracket"),
      Token::Comma => write!(f, "Comma"),
      Token::Eof => write!(f, "Eof"),
    }
  }
}

// 词法分析器
struct Lexer<'a> {
  chars: Chars<'a>,
  current: Option<char>,
}

impl<'a> Lexer<'a> {
  fn new(input: &'a str) -> Self {
    let mut chars = input.chars();
    let current = chars.next();
    Lexer { chars, current }
  }

  // 前进到下一个字符
  fn advance(&mut self) {
    self.current = self.chars.next();
  }

  // 跳过空白字符
  fn skip_whitespace(&mut self) {
    while let Some(c) = self.current {
      if c.is_whitespace() {
        self.advance();
      } else {
        break;
      }
    }
  }

  // 解析数字
  fn parse_number(&mut self) -> f64 {
    let mut num_str = String::new();
    while let Some(c) = self.current {
      if c.is_ascii_digit() || c == '.' {
        num_str.push(c);
        self.advance();
      } else {
        break;
      }
    }
    num_str.parse().unwrap_or(0.0)
  }

  // 获取下一个令牌
  fn next_token(&mut self) -> Token {
    self.skip_whitespace();

    match self.current {
      None => Token::Eof,
      Some('+') => {
        self.advance();
        Token::Plus
      }
      Some('-') => {
        self.advance();
        Token::Minus
      }
      Some('*') => {
        self.advance();
        Token::Star
      }
      Some('/') => {
        self.advance();
        Token::Slash
      }
      Some('(') => {
        self.advance();
        Token::LParen
      }
      Some(')') => {
        self.advance();
        Token::RParen
      }
      Some('[') => {
        // 新增数组左括号处理
        self.advance();
        Token::LBracket
      }
      Some(']') => {
        // 新增数组右括号处理
        self.advance();
        Token::RBracket
      }
      Some(',') => {
        // 新增逗号处理
        self.advance();
        Token::Comma
      }
      Some(c) if c.is_ascii_digit() => Token::Number(self.parse_number()),
      Some(c) => {
        panic!("Unexpected character: {}", c);
      }
    }
  }
}

// 表达式节点
#[derive(Debug, PartialEq)]
enum Expr {
  Number(f64),
  Binary {
    op: Token,
    left: Box<Expr>,
    right: Box<Expr>,
  },
  Unary {
    op: Token,
    expr: Box<Expr>,
  },
  Array(Vec<Expr>), // 数组节点，包含表达式列表
}

// Pratt解析器
struct Parser<'a> {
  lexer: Lexer<'a>,
  current_token: Token,
}

impl<'a> Parser<'a> {
  fn new(lexer: Lexer<'a>) -> Self {
    let mut parser = Parser {
      lexer,
      current_token: Token::Eof,
    };
    parser.advance();
    parser
  }

  // 前进到下一个令牌
  fn advance(&mut self) {
    self.current_token = self.lexer.next_token();
  }

  // 获取当前令牌的优先级
  fn precedence(token: &Token) -> u8 {
    match token {
      Token::Plus | Token::Minus => 1,
      Token::Star | Token::Slash => 2,
      _ => 0,
    }
  }

  // 解析数组：处理 [expr, expr, ...] 格式
  fn parse_array(&mut self) -> Expr {
    self.advance(); // 跳过左括号 [
    let mut elements = Vec::new();

    // 处理空数组 []
    if let Token::RBracket = self.current_token {
      self.advance();
      return Expr::Array(elements);
    }

    // 解析数组元素
    loop {
      // 数组元素可以是任意表达式
      let elem = self.parse_expr(0);
      elements.push(elem);

      // 检查下一个令牌是逗号还是右括号
      match &self.current_token {
        Token::Comma => {
          self.advance(); // 跳过逗号
        }
        Token::RBracket => {
          self.advance(); // 跳过右括号
          break;
        }
        token => panic!("Unexpected token in array: {}", token),
      }
    }

    Expr::Array(elements)
  }
  // 解析原子表达式（数字、括号表达式、一元运算符）
  fn parse_atom(&mut self) -> Expr {
    match &self.current_token {
      Token::Number(n) => {
        let val = *n;
        self.advance();
        Expr::Number(val)
      }
      Token::LParen => {
        self.advance(); // 跳过 '('
        let expr = self.parse_expr(0);
        if let Token::RParen = self.current_token {
          self.advance(); // 跳过 ')'
        } else {
          panic!("Expected ')'");
        }
        expr
      }
      Token::LBracket => {
        // 新增数组解析入口
        self.parse_array()
      }
      Token::Plus | Token::Minus => {
        // 一元运算符
        let op = self.current_token.clone();
        self.advance();
        let rhs = self.parse_atom();
        Expr::Unary {
          op,
          expr: Box::new(rhs),
        }
      }
      token => panic!("Unexpected token in atom: {}", token),
    }
  }

  // 解析表达式的核心方法
  fn parse_expr(&mut self, min_prec: u8) -> Expr {
    let mut left = self.parse_atom();

    while min_prec < Self::precedence(&self.current_token) {
      let op = self.current_token.clone();
      self.advance();

      let rhs = self.parse_expr(Self::precedence(&op));
      left = Expr::Binary {
        op,
        left: Box::new(left),
        right: Box::new(rhs),
      };
    }

    left
  }

  // 解析整个表达式
  fn parse(&mut self) -> Expr {
    let expr = self.parse_expr(0);
    if let Token::Eof = self.current_token {
      expr
    } else {
      panic!("Unexpected token at end: {}", self.current_token);
    }
  }
}

// // 计算表达式的值
// fn evaluate(expr: &Expr) -> f64 {
//   match expr {
//     Expr::Number(n) => *n,
//     Expr::Unary { op, expr } => {
//       let val = evaluate(expr);
//       match op {
//         Token::Plus => val,
//         Token::Minus => -val,
//         _ => panic!("Invalid unary operator"),
//       }
//     }
//     Expr::Binary { op, left, right } => {
//       let left_val = evaluate(left);
//       let right_val = evaluate(right);
//       match op {
//         Token::Plus => left_val + right_val,
//         Token::Minus => left_val - right_val,
//         Token::Star => left_val * right_val,
//         Token::Slash => left_val / right_val,
//         _ => panic!("Invalid binary operator"),
//       }
//     }
//   }
// }

// 扩展表达式计算，支持数组求值
fn evaluate(expr: &Expr) -> Vec<f64> {
  match expr {
    Expr::Number(n) => vec![*n], // 单个数字作为单元素数组返回
    Expr::Unary { op, expr } => {
      let val = evaluate(expr);
      assert_eq!(val.len(), 1, "Unary operator requires single value");
      match op {
        Token::Plus => val,
        Token::Minus => vec![-val[0]],
        _ => panic!("Invalid unary operator"),
      }
    }
    Expr::Binary { op, left, right } => {
      let left_vals = evaluate(left);
      let right_vals = evaluate(right);
      assert_eq!(
        left_vals.len(),
        1,
        "Binary left operand must be single value"
      );
      assert_eq!(
        right_vals.len(),
        1,
        "Binary right operand must be single value"
      );

      let left_val = left_vals[0];
      let right_val = right_vals[0];

      match op {
        Token::Plus => vec![left_val + right_val],
        Token::Minus => vec![left_val - right_val],
        Token::Star => vec![left_val * right_val],
        Token::Slash => vec![left_val / right_val],
        _ => panic!("Invalid binary operator"),
      }
    }
    Expr::Array(elements) => {
      // 数组求值：计算每个元素并返回结果数组
      elements
        .iter()
        .flat_map(evaluate) // 展平每个元素的计算结果（单个值）
        .collect()
    }
  }
}
fn main() {
  let tests = [
    ("3 + 4 * 2 / (1 - 5)", -2.5),
    ("1 + 2 * 3", 7.0),
    ("(5 + 3) * 2", 16.0),
    ("-10 + 5", -5.0),
    ("10 / 2 - 3", 2.0),
    ("2 + -3 * 4", -10.0),
  ];

  for (input, expected) in tests {
    let lexer = Lexer::new(input);
    let mut parser = Parser::new(lexer);
    let expr = parser.parse();
    let result = evaluate(&expr);
    println!(
      "Input: '{}' => Parsed: {:?} => Result: {:.2?} (Expected: {:.2})",
      input, expr, result, expected
    );
  }

  let tests = [
    ("[1, 3]", vec![1.0, 3.0]),
    ("[1 + 2, 4 * 5]", vec![3.0, 20.0]),
    ("[ (5 - 3), 10 / 2 ]", vec![2.0, 5.0]),
    ("[ -1, +4 ]", vec![-1.0, 4.0]),
    ("[1 + 2 * 3, (4 + 5) / 3]", vec![7.0, 3.0]),
    ("[]", vec![]), // 测试空数组
  ];

  for (input, expected) in tests {
    let lexer = Lexer::new(input);
    let mut parser = Parser::new(lexer);
    let expr = parser.parse();
    let result = evaluate(&expr);
    println!(
      "Input: '{}' => Parsed: {:?} => Result: {:?} (Expected: {:?})",
      input, expr, result, expected
    );
    assert_eq!(result, expected);
  }
}
